#include <assert.h>
#include <stdlib.h>

// CUDA runtime
#include "helper.h"
#include <cublas_v2.h>
#include <cuda_runtime.h>

#define ROW_MAJOR(row, col, ld) ((row) * (ld) + (col))
#define FLOAT4(pointer) (reinterpret_cast<float4*>(&(pointer))[0])
#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))

constexpr int kWarpSize = 32;

constexpr int BM = 64;
constexpr int BN = 128;
constexpr int BK = 16;
constexpr int TM = BM / 8;
constexpr int TN = BN / 16;
constexpr int kBlockDim = (BM / TM) * (BN / TN);
static_assert(kBlockDim >= 4 * kWarpSize, "At least 4 warps per block are required for performance");
static_assert(BM * BK >= 4 * kBlockDim, "Each thread needs at least 4 float to load A");
static_assert(BN * BK >= 4 * kBlockDim, "Each thread needs at least 4 float to load B");

// Global A shape:  [(m/(BM), BM), (k/BK, BK/4, 4)]
// Global A stride: [(  k*BM, k ), (BK  , 4 , 1)]
// Global A index:      bx    tx    bk  tid%bk  

// Align4 for a, b, c to use float4 load/store
__global__ void sgemm_align4(const int M, const int N, const int K,
    float * __restrict__ a, float * __restrict__ b, float * __restrict__ c) {

    const int bx = blockIdx.x;
    const int by = blockIdx.y;
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;

    __shared__ float s_a[BM][BK];
    __shared__ float s_b[BK][BN];

    float r_c[TM][TN] = {0.0};

    constexpr int kThreadPerBK = BK / 4; // Align4
    constexpr int kThreadPerBN = BN / 4; // Align4
    constexpr int kLoadAIterM = CEIL_DIV(BM * BK, 4 * kBlockDim);  // Number of iterations required to load one block of A
    constexpr int kLoadBIterK = CEIL_DIV(BK * BN, 4 * kBlockDim);  // Number of iterations required to load one block of B
    // constexpr int kStoreCIterM = CEIL_DIV(BM * BN, 4 * kBlockDim);  // Number of iterations required to store one block of C
    constexpr int kLoadAStepM = CEIL_DIV(4 * kBlockDim, BK);  // Step for a thread for loading one block of A from global to shared on M dimension
    constexpr int kLoadBStepK = CEIL_DIV(4 * kBlockDim, BN);  // Step of iterations required to load one block of B
    // constexpr int kStoreCStepM = CEIL_DIV(4 * kBlockDim, BN);  // Number of iterations required to store one block of C

    int load_a_smem_m =  tid / kThreadPerBK;
    int load_a_smem_k = (tid  %  kThreadPerBK) * 4;
    int load_b_smem_k =  tid / kThreadPerBN;
    int load_b_smem_n = (tid % kThreadPerBN) * 4;

    int load_a_gmem_m = by * BM + load_a_smem_m;
    int load_b_gmem_n = bx * BN + load_b_smem_n;

    for (int bk = 0; bk < CEIL_DIV(K, BK); bk++) {
        int load_a_gmem_k = bk * BK + load_a_smem_k;
        int load_b_gmem_k = bk * BK + load_b_smem_k;
        #pragma unroll
        for (int iterM = 0; iterM < kLoadAIterM; iterM++) {
            int offset = iterM * kLoadAStepM;
            int load_a_gmem_addr = ROW_MAJOR(load_a_gmem_m + offset, load_a_gmem_k, K);
            FLOAT4(s_a[load_a_smem_m + offset][load_a_smem_k]) = FLOAT4(a[load_a_gmem_addr]);
        }
        #pragma unroll
        for (int iterK = 0; iterK < kLoadBIterK; iterK++) {
            int offset = iterK * kLoadBStepK;
            int load_b_gmem_addr = ROW_MAJOR(load_b_gmem_k + offset, load_b_gmem_n, N);
            FLOAT4(s_b[load_b_smem_k + offset][load_b_smem_n]) = FLOAT4(b[load_b_gmem_addr]);
        }

        __syncthreads();

        #pragma unroll
        for (int m = 0; m < TM; m++) {
            #pragma unroll
            for (int n = 0; n < TN; n++) {
                #pragma unroll
                for (int k = 0; k < BK; k++) {
                    int comp_a_smem_m = ty * TM + m;
                    int comp_b_smem_n = tx * TN + n;
                    r_c[m][n] += s_a[comp_a_smem_m][k] * s_b[k][comp_b_smem_n];
                }
            }
        }

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < TM; i++) {
        int store_c_gmem_m = by * BM + ty * TM + i;
        #pragma unroll
        for (int j = 0; j < TN; j += 4) {
            int store_c_gmem_n = bx * BN + tx * TN + j;
            int store_c_gmem_addr = ROW_MAJOR(store_c_gmem_m, store_c_gmem_n, N);
            FLOAT4(c[store_c_gmem_addr]) = FLOAT4(r_c[i][j]);
        }
    }
}



void SGEMM_GPU(cublasHandle_t handle, int m, int n, int k, float *d_A, float *d_B, float *d_C) {

  dim3 block(BN/TN, BM/TM);
  dim3 grid(ceil(n/BN), ceil(m/BM));  

  sgemm_align4<<<grid, block>>>(m, n, k, d_A, d_B, d_C);
}
